from pprint import pprint

from paddlenlp import Taskflow

# 实体抽取任务
schema = ["时间", "选手", "赛事名称", "犯罪嫌疑人"]  # Define the schema for entity extraction
ie = Taskflow(
    task="information_extraction",
    model=r"D:\uie",
    schema=schema,
    batch_size=3, # 批处理大小，处理多个输入样本时，可以加速处理速度，但会增加内存消耗。默认值为1
    precision="float16",
)
pprint(
    ie(
        [
            "被告人朱某某于2018年4月25日至4月27日期间，先后在常州市钟楼区**镇**东大门**路**号门口、**村委**村**号**楼**房间、**路**物流港**幢*号**楼一房间等地盗窃作案3次，窃得被害人吴某某、周某某、王某某的电动车、手机等物品。",
            "2月12日，在哈尔滨亚冬会越野滑雪男子4×7.5公里接力决赛中，由李明林、次仁占堆、宝林、王强组成的中国队夺得金牌。",
            "2015年2月6日公司接到股东深圳市傲盛霞实业有限公司（以下简称“傲盛霞”）的通知：傲盛霞于2015年2月6日收到广东省深圳市中级人民法院的解除查封冻结通知书，广东省深圳市中级人民法院于2015年2月5日解除傲盛霞持有的我司18220000股（证券代码：000622，简称：恒立实业）的冻结。",
        ]
    )
)